/**
 * @file general tools/functions for computations
 *
 **/


#include <math.h>
#include <NTL/ZZ_pXFactoring.h>
#include <NTL/ZZ_pEX.h>
#include <NTL/matrix.h>
#include <NTL/vec_vec_ZZ_pE.h>
#include <NTL/mat_ZZ_pE.h>
#include <NTL/mat_GF2.h>
#include <NTL/mat_ZZ_p.h>
#include <NTL/ZZ_pEXFactoring.h>

#include "SQ_ring.h"

NTL_CLIENT

/**
 * Computes combination n choose k
 *
 */
unsigned long long nChoosek(unsigned n, unsigned k) 
{
    if (k > n)
        return 0;

    if (k > n/2)
        k = n-k; // Take advantage of symmetry

    long double accum = 1;
    unsigned i;
    for (i = 1; i <= k; i++)
         accum = accum * (n-k+i) / i;

    return accum + 0.5; // avoid rounding error
}

/**
 * Given a skew generator polynomial, returns the generator matrix
 *
 * @gen_poly	skew generator polynomial
 * @n	code length
 *
 * @return	generator matrix
 **/
mat_ZZ_pE SQ_gen_matrix(const ZZ_pEX& gen_poly,int n)
{
	ZZ_pEX gen_polyCPY=gen_poly;
	ZZ_pEX pEX_X;
	SetX(pEX_X); // pEX_X = x
	int r = deg(gen_poly);
	vec_vec_ZZ_pE gen_matrix_vec;
	append(gen_matrix_vec,VectorCopy(gen_polyCPY,n));
	for(int i=1;i<n-r;i++)	// multiply gen_poly by x from the left each iteration and assign to a row of the generator matrix
	{
		gen_polyCPY=SQ_mulr(pEX_X,gen_polyCPY);
		append(gen_matrix_vec,VectorCopy(gen_polyCPY,n));
	}
	mat_ZZ_pE gen_matrix;
	MakeMatrix(gen_matrix,gen_matrix_vec);
	return gen_matrix;
}

/**
 * Generates a list of all messages of length k with weight j in GF(2)
 *
 * @list list of vectors, Note: memory must be preallocated for the list of size [k choose j,k]
 * @k	length of messages 
 * @j	weight
 *
 * @return size of the list
 *
 **/
int binary_msg_generator(vec_ZZ_pE* list,int k,int j,int horizontal_index=0,int vertical_index=0)
{
	ZZ_p p_ONE;
	set(p_ONE); // p_ONE = 1
	int size = list[0].length();
	int kChoosej= nChoosek(k,j);
	int vertical_stop_index = vertical_index+kChoosej;
	if( kChoosej ==1 ) 
	{
		if( k == j )
		{
			for(int i=horizontal_index;i<size;i++)
				list[vertical_index][i]= p_ONE; // set element to 1
		}
		return(vertical_index+1);
	}
	while(vertical_index<vertical_stop_index)
	{
		int updated_vertical_index = binary_msg_generator(list,k-1,j-1,horizontal_index+1,vertical_index);
		for(;vertical_index<updated_vertical_index;vertical_index++)
			list[vertical_index][horizontal_index]=p_ONE;
		horizontal_index++;	
		k--;
	}
	return vertical_index;
}

/**
 *  A counter that increments the input vector in binary by 1.
 *
 *  In other words this is the implementation of a digital full adder. For example, if input is [0 1 0 1] it outputs [0 1 1 0].  
 *
 **/
vec_ZZ_p binary_inc(vec_ZZ_p vec)
{
	int n=vec.length();
	vec_ZZ_p res;
	res.SetLength(n);
	ZZ_p c_in,p_ONE;
	set(p_ONE); // p_ONE = 1
	set(c_in); // c_in = 1
	for(int i=0; i<n;i++)
	{
		res[i] = vec[i]*(c_in+p_ONE)+(vec[i]+p_ONE)*c_in;
		c_in = vec[i]*c_in;	
	}
	return res;
}
